package com.nanxhs.structure.tree;

import java.util.Objects;

/**
 *
 * 二叉查找平衡树
 *
 * @author 唐海斌
 * @version 1.0
 * @date 2021/5/25 9:58
 */
public class AVLTree extends Tree {

    @Override
    public void add(int element) {
        System.out.println("add element: " + element);
        TreeNode newNode = new TreeNode();
        newNode.setElement(element);
        if (root == null) {
            root = newNode;
            return;
        }
        add(root, newNode);
    }

    private void add(TreeNode parentNode, TreeNode newNode) {
        TreeNode childrenLeft = parentNode.getChildrenLeft();
        TreeNode childrenRight = parentNode.getChildrenRight();
        Integer parentElement = parentNode.getElement();
        Integer newElement = newNode.getElement();
        if (Objects.equals(parentElement, newElement)) {
            return;
        }
        if (newElement < parentElement) {
            if (childrenLeft == null) {
                newNode.setParent(parentNode);
                parentNode.setChildrenLeft(newNode);
                updateBalanceFactor(newNode);
                return;
            }
            add(childrenLeft, newNode);
        } else {
            if (childrenRight == null) {
                newNode.setParent(parentNode);
                parentNode.setChildrenRight(newNode);
                updateBalanceFactor(newNode);
                return;
            }
            add(childrenRight, newNode);
        }
    }

    /**
     * 旋转
     * @param node
     */
    private void rotate(TreeNode node) {
        TreeNode parentNode = node.getParent();
        if (parentNode == null) {
            return;
        }
        NodeLocation parentNodeLocation = getNodeLocation(parentNode);
        NodeLocation nodeLocation = getNodeLocation(node);
        //如果当前节点处于父节点的左边
        if (nodeLocation == NodeLocation.left) {
            //如果父节点在左边或者父节点为root节点  LL情况 右旋
            if (parentNodeLocation == NodeLocation.left || parentNodeLocation == NodeLocation.root) {
                rightRotate(node);
                return;
            }
            //LR情况, 先左旋后右旋
            leftRightRotate(node);
            return;
        }
        //如果当前节点处于父节点的右边
        if (nodeLocation == NodeLocation.right) {
            //如果父节点在右边或者父节点为root节点  RR情况 左旋
            if (parentNodeLocation == NodeLocation.right || parentNodeLocation == NodeLocation.root) {
                leftRotate(node);
                return;
            }
            //RL情况, 先右旋后左旋
            rightLeftRotate(node);
        }
    }

    /**
     * 更新平衡因子
     * @param node
     */
    private void updateBalanceFactor(TreeNode node) {
        TreeNode parentNode = node.getParent();
        if (parentNode == null) {
            return;
        }
        NodeLocation nodeLocation = getNodeLocation(node);
        if (nodeLocation == NodeLocation.left) {
            parentNode.setBalanceFactor(parentNode.getBalanceFactor() + 1);
        }
        if (nodeLocation == NodeLocation.right) {
            parentNode.setBalanceFactor(parentNode.getBalanceFactor() - 1);
        }
        TreeNode grandfatherNode = parentNode.getParent();
        if (grandfatherNode == null) {
            return;
        }
        int balanceFactor = grandfatherNode.getBalanceFactor();
        if (nodeLocation == NodeLocation.left) {
            balanceFactor = balanceFactor + 1;
            if (Math.abs(balanceFactor) > 1) {
                rotate(node);
            }
//            updateBalanceFactor(parentNode);
            return;
        }
        if (nodeLocation == NodeLocation.right) {
            balanceFactor = balanceFactor - 1;
            if (Math.abs(balanceFactor) > 1) {
                rotate(node);
            }
//            updateBalanceFactor(parentNode);
        }
    }

    /**
     * 左子树高度 - 右子树高度 的绝对值 小于2 说明目前是平衡树
     *
     * @param node
     * @return
     */
    private boolean isBalanceFactorTree(TreeNode node) {
        TreeNode parentNode = node.getParent();
        //父节点为空 说明当前节点是root节点 那么 子树高度为1
        if (parentNode == null) {
            return true;
        }
        TreeNode grandfatherNode = parentNode.getParent();
        //说明是根节点，那么当前node的爷爷子树高为2，所以一定是平衡树
        if (grandfatherNode == null) {
            return true;
        }
        NodeLocation parentNodeLocation = getNodeLocation(parentNode);
        //父节点处于爷爷节点的左边
        if (parentNodeLocation == NodeLocation.left) {
            return Math.abs(2 - getNodeHeight(grandfatherNode.getChildrenRight(), 0)) < 2;
        }
        //父节点处于爷爷节点的右边
        return Math.abs(getNodeHeight(grandfatherNode.getChildrenLeft(), 0) - 2) < 2;
    }

    /**
     * 获取节点的左树高度
     * @return
     */
    private int getNodeHeight(TreeNode node, int height) {
        if (node == null) {
            return height;
        }
        TreeNode childrenLeft = node.getChildrenLeft();
        TreeNode childrenRight = node.getChildrenRight();
        if (childrenLeft == null && childrenRight == null) {
            return height + 1;
        }
        if (childrenLeft != null && childrenRight != null) {
            int nodeLeftHeight = getNodeHeight(childrenLeft, height + 1);
            int nodeRightHeight = getNodeHeight(childrenRight, height + 1);
            return Math.max(nodeLeftHeight, nodeRightHeight);
        }
        if (childrenLeft != null && childrenRight == null) {
            return getNodeHeight(childrenLeft, height + 1);
        }
        return getNodeHeight(childrenRight, height + 1);
    }

    private NodeLocation getNodeLocation(TreeNode node) {
        if (node == root) {
            return NodeLocation.root;
        }
        return node.getParent().getChildrenLeft() == node ? NodeLocation.left : NodeLocation.right;
    }

    /**
     * 左旋
     *
     * 右子树出现失衡并且其祖先节点处于右节点
     *
     * @param node 失衡节点
     */
    private void leftRotate(TreeNode node) {
        //失衡节点的父节点
        TreeNode unbalanceParent = node.getParent();
        unbalanceParent.setBalanceFactor();
        //失衡节点的爷爷
        TreeNode grandfatherNode = unbalanceParent.getParent();
        if (grandfatherNode == null) {
            root = node;
        } else {
            grandfatherNode.setChildrenRight(node);
        }
        //失衡节点占用父节点位置
        node.setParent(grandfatherNode);
        //失衡节点的父亲成为失衡节点的左子节点
        unbalanceParent.setParent(node);
        node.setChildrenLeft(unbalanceParent);
        unbalanceParent.setChildrenRight(null);
    }

    /**
     * 右旋
     *
     * 左子树出现失衡并且其祖先节点处于左节点
     * @param node 失衡节点
     */
    private void rightRotate(TreeNode node) {
        //失衡节点的父节点
        TreeNode unbalanceParent = node.getParent();
        //失衡节点的爷爷
        TreeNode grandfatherNode = unbalanceParent.getParent();
        if (grandfatherNode == null) {
            root = node;
        } else {
            grandfatherNode.setChildrenLeft(node);
        }
        //失衡节点占用父节点位置
        node.setParent(grandfatherNode);
        //失衡节点的父亲成为失衡节点的左子节点
        unbalanceParent.setParent(node);
        node.setChildrenRight(unbalanceParent);
        unbalanceParent.setChildrenLeft(null);
    }

    /**
     * 先左旋后右旋
     *
     * 右子树出现失衡并且祖先节点处于左节点
     * @param node
     */
    private void leftRightRotate(TreeNode node) {
        leftRotate(node);
        rightRotate(node);
    }

    /**
     * 先右旋后左旋
     *
     * 左子树出现失衡并且祖先节点处于右节点
     * @param node
     */
    private void rightLeftRotate(TreeNode node) {
        rightRotate(node);
        leftRotate(node);
    }

    @Override
    void remove(int element) {

    }

    /**
     *  节点位置
     */
    enum NodeLocation {
        /**
         * 根节点
         */
        root,
        /**
         * 左节点
         */
        left,
        /**
         * 右节点
         */
        right
        ;
    }
}
